# Copyright (c) Meta Platforms, Inc. and affiliates.
# This software may be used and distributed according to the terms of the Llama 2 Community License Agreement.

# For dataset details visit: https://crfm.stanford.edu/2023/03/13/alpaca.html

import copy

import torch
from torch.utils.data import Dataset

from .utils import JsonlLoader


PROMPT_TEMPLATE_DICT = {
    "prompt_template_en": (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n"
        "### Instruction:\n"
        "{query}\n"
        "### Response:"
    ),
    "prompt_template_cn": (
        "你是一个中医临床医生，用正确的治疗方法来治疗疾病。"
        "你能够推荐常规药物、草药和其他天然替代品在提供建议时，还需要考虑患者的年龄、生活方式和病史。\n"
        "我的第一个建议请求是“{query}”。\n"
        "建议: "
    ),
}


class InstructionDataset(Dataset):
    def __init__(self, dataset_config, tokenizer, partition="train", max_tokens=256):
        self.data = JsonlLoader(dataset_config.data_path, True)
        if partition == "train":
            self.data = self.data
        elif partition == "val":
            self.data = self.data[:300]
        else:
            raise NotImplementedError(f'Partition {partition} is not availiable.')

        self.max_tokens = max_tokens
        # tokenizer = Tokenizer(model_path=model_path + "./tokenizer.model")
        self.tokenizer = tokenizer
        # self.tokenizer1 = tokenizer

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        IGNORE_INDEX = -100  # The default setting in CrossEntropyLoss
        sample_from_data = self.data[index]
        prompt = PROMPT_TEMPLATE_DICT["prompt_template_cn"].format_map(sample_from_data)
        prompt_ids = torch.tensor(self.tokenizer.encode(prompt), dtype=torch.int64)

        example = prompt + sample_from_data["response"]
        example_ids = self.tokenizer.encode(example)
        example_ids.append(self.tokenizer.eos_token_id)
        example_ids = torch.tensor(example_ids, dtype=torch.int64)

        padding_num = self.max_tokens - example_ids.shape[0]
        if padding_num > 0:
            example_ids = torch.cat((example_ids, torch.zeros(padding_num, dtype=torch.int64) - 1))
        elif padding_num < 0:
            example_ids = example_ids[:self.max_tokens]
        labels = copy.deepcopy(example_ids)
        labels[: len(prompt_ids)] = -1
        example_mask = example_ids.ge(0)
        label_mask = labels.ge(0)
        example_ids[~example_mask] = 0
        labels[~label_mask] = IGNORE_INDEX
        example_mask = example_mask.float()
        label_mask = label_mask.float()

        return {
            "input_ids": example_ids,
            "labels": labels,
            "attention_mask": example_mask,
        }
